"""
数据库迁移脚本：添加学院和系部字段
为users表添加college和department字段，并将email改为可选
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dotenv import load_dotenv
load_dotenv()

from app import create_app
from app.extensions import db


def migrate_database():
    """执行数据库迁移"""
    app = create_app('development')
    
    with app.app_context():
        print("=" * 60)
        print("开始数据库迁移：添加学院和系部字段")
        print("=" * 60)
        
        try:
            with db.engine.begin() as conn:
                # 检查字段是否已存在
                result = conn.execute(db.text("""
                    SELECT COUNT(*) as count 
                    FROM information_schema.COLUMNS 
                    WHERE TABLE_SCHEMA = DATABASE() 
                    AND TABLE_NAME = 'users' 
                    AND COLUMN_NAME = 'college'
                """))
                exists = result.fetchone()[0] > 0
                
                if exists:
                    print("\n⚠️  字段已存在，跳过迁移")
                    return
                
                print("\n添加college字段...")
                conn.execute(db.text("""
                    ALTER TABLE users 
                    ADD COLUMN college VARCHAR(100) DEFAULT '' NOT NULL
                """))
                print("✓ college字段添加成功")
                
                print("\n添加department字段...")
                conn.execute(db.text("""
                    ALTER TABLE users 
                    ADD COLUMN department VARCHAR(100) DEFAULT '' NOT NULL
                """))
                print("✓ department字段添加成功")
                
                print("\n修改email字段为可选...")
                conn.execute(db.text("""
                    ALTER TABLE users 
                    MODIFY COLUMN email VARCHAR(100) NULL
                """))
                print("✓ email字段修改成功")
                
                print("\n修改first_name字段为必填...")
                conn.execute(db.text("""
                    ALTER TABLE users 
                    MODIFY COLUMN first_name VARCHAR(50) NOT NULL
                """))
                print("✓ first_name字段修改成功")
                
                print("\n修改last_name字段为必填...")
                conn.execute(db.text("""
                    ALTER TABLE users 
                    MODIFY COLUMN last_name VARCHAR(50) NOT NULL
                """))
                print("✓ last_name字段修改成功")
                
                print("\n更新现有用户数据...")
                # 为现有用户设置默认学院和系部
                conn.execute(db.text("""
                    UPDATE users 
                    SET college = '未分配学院', 
                        department = '未分配系部'
                    WHERE college = '' OR department = ''
                """))
                print("✓ 现有用户数据更新成功")
                
            print("\n" + "=" * 60)
            print("数据库迁移完成！")
            print("=" * 60)
            
        except Exception as e:
            print(f"\n❌ 迁移失败: {str(e)}")
            print("\n提示：如果是全新数据库，请直接运行 init_database.py")
            raise


if __name__ == '__main__':
    migrate_database()


